from enum import Enum
from typing import TypedDict,Any,List,Dict,Union,cast
from soj.common.judge import JudgeResultCode
class ScorePolicy(Enum):
    SUM="SUM"
    MIN="MIN"
    MAX="MAX"
class TimePolicy(Enum):
    SUM="SUM"
    MIN="MIN"
    MAX="MAX"
class MemoryPolicy(Enum):
    SUM="SUM"
    MIN="MIN"
    MAX="MAX"

class Policy:
    def __init__(self,sp:ScorePolicy=ScorePolicy.MIN,
                    tp:TimePolicy=TimePolicy.MAX,
                    mp:MemoryPolicy=MemoryPolicy.MAX):
        self.score_policy=sp
        self.time_policy=tp
        self.memory_policy=mp

def get_policy_by_dict(d:dict)->Policy:
    try:
        if d["score"]=="SUM":
            sp=ScorePolicy.SUM
        elif d["score"]=="MAX":
            sp=ScorePolicy.MAX
        elif d["score"]=="MIN":
            sp=ScorePolicy.MIN
        else:
            raise ValueError("Invalid score policy")
        if d["time"]=="SUM":
            tp=TimePolicy.SUM
        elif d["time"]=="MAX":
            tp=TimePolicy.MAX
        elif d["time"]=="MIN":
            tp=TimePolicy.MIN
        else:
            raise ValueError("Invalid time policy")
        if d["memory"]=="SUM":
            mp=MemoryPolicy.SUM
        elif d["memory"]=="MAX":
            mp=MemoryPolicy.MAX
        elif d["memory"]=="MIN":
            mp=MemoryPolicy.MIN
        else:
            raise ValueError("Invalid memory policy")
        return Policy(sp,tp,mp)
    except KeyError as e:
        raise ValueError("Incomplete policy dict") from e

class ScorableObject:
    def __init__(self):
        pass
    def to_json(self,paintable=True)->Any:
        return {}
    def get_score(self):
        return 0
    def get_time(self):
        return 0
    def get_memory(self):
        return 0
    def get_status(self)->JudgeResultCode:
        return JudgeResultCode.SYSTEM_ERROR

class Testcase(ScorableObject):
    def __init__(self,id:int,score:int,time:int=0,mem:int=0,status:JudgeResultCode|None=None):
        self.score=score
        self.time=time
        self.memory=mem
        self.id=id
        self.status=status
    def to_json(self,paintable=True)->Any:
        result={
            'id':self.id,
            'score':self.score,
            'time':self.time,
            'memory':self.memory,
            'status':self.status if self.status else JudgeResultCode.SYSTEM_ERROR
        }
        if paintable:result["status"]=result["status"].value
        return result
    def get_score(self):
        return self.score
    def get_time(self):
        return self.time
    def get_memory(self):
        return self.memory
    def get_status(self):
        return self.status
    

class SubtaskResult(TypedDict):
    score:int
    time:int
    memory:int
    status:JudgeResultCode



class Subtask(ScorableObject):
    def __init__(self,id:int,policy:Policy=Policy()):
        self.policy=policy
        self.id=id
        self.cases:List[ScorableObject]=[]
        self.status:JudgeResultCode=JudgeResultCode.SYSTEM_ERROR
    def add_testcase(self,result:int,time:int=0,mem:int=0,id:int|None=None):
        self.cases.append(Testcase(len(self.cases) if not id else id,result,time,mem))
        return self.cases[-1]
    def add_subtask(self,pol:Policy=Policy(),id:int|None=None):
        self.cases.append(Subtask(len(self.cases) if not id else id,pol))
        return self.cases[-1]
    def get_score(self):
        if self.policy.score_policy==ScorePolicy.SUM:
            return sum(map(lambda x:x.get_score(),self.cases))
        elif self.policy.score_policy==ScorePolicy.MAX:
            return max(map(lambda x:x.get_score(),self.cases))
        elif self.policy.score_policy==ScorePolicy.MIN:
            return min(map(lambda x:x.get_score(),self.cases))
        else:
            raise NotImplementedError("Not supported scoring policy")
    def get_time(self):
        if self.policy.time_policy==TimePolicy.SUM:
            return sum(map(lambda x:x.get_time(),self.cases))
        elif self.policy.time_policy==TimePolicy.MAX:
            return max(map(lambda x:x.get_time(),self.cases))
        elif self.policy.time_policy==TimePolicy.MIN:
            return min(map(lambda x:x.get_time(),self.cases))
        else:
            raise NotImplementedError("Not supported timing policy")
    def get_memory(self):
        if self.policy.memory_policy==MemoryPolicy.SUM:
            return sum(map(lambda x:x.get_memory(),self.cases))
        elif self.policy.memory_policy==MemoryPolicy.MAX:
            return max(map(lambda x:x.get_memory(),self.cases))
        elif self.policy.memory_policy==MemoryPolicy.MIN:
            return min(map(lambda x:x.get_memory(),self.cases))
        else:
            raise NotImplementedError("Not supported memory counting policy")
    def get_status(self):
        statuses=list(map(lambda x:x.get_status(),self.cases))
        
        if JudgeResultCode.SYSTEM_ERROR in statuses:
            return JudgeResultCode.SYSTEM_ERROR
        if JudgeResultCode.RUNTIME_ERROR in statuses:
            return JudgeResultCode.UNACCEPTABLE
        if JudgeResultCode.MEMORY_LIMIT_EXCEEDED in statuses:
            return JudgeResultCode.UNACCEPTABLE
        if JudgeResultCode.TIME_LIMIT_EXCEEDED in statuses:
            return JudgeResultCode.UNACCEPTABLE
        if JudgeResultCode.WRONG_ANSWER in statuses:
            return JudgeResultCode.UNACCEPTABLE
        if JudgeResultCode.UNACCEPTABLE in statuses:
            return JudgeResultCode.UNACCEPTABLE
        return JudgeResultCode.ACCEPTABLE
    def get_result(self)->SubtaskResult:
        return {
            'score':self.get_score(),
            'time':self.get_time(),
            'memory':self.get_memory(),
            'status':self.get_status()
        }
    def to_json(self,paintable=True)->Any:
        res=[]
        for i in self.cases:
            res.append(i.to_json(paintable))
        dct={
            'id':self.id,
            'result':res
        }
        dct.update(self.get_result())
        if paintable:dct["status"]=dct["status"].value
        return dct

class TestcaseConfig(TypedDict):
    real_id:int

class SubtaskConfig(TypedDict):
    policy:Policy
    real_id:int
    sub:List["TestcaseConfig|SubtaskConfig"]

class InvalidConfigError(RuntimeError):
    def __init__(self):
        super().__init__("Invalid scoring config")

class ScoringEngine:
    
    def __init__(self,conf:SubtaskConfig):
        self.usub:Subtask|None=None
        self.testcase_map:Dict[int,List[Testcase]]={}
        self.subtask_map:Dict[int,Subtask]={}
        self.root_id:int=-1    
        try:
            conf["real_id"]=self.root_id #root id must be -1
            self.usub=Subtask(self.root_id,conf["policy"])
            self.subtask_map[self.root_id]=self.usub
            self._build_engine(self.usub,conf)
        except Exception as e:
            raise InvalidConfigError() from e
    def get_result(self)->SubtaskResult:
        return self.usub.get_result()
    def to_json(self,paintable=True)->Any:
        return self.usub.to_json()
    def set_testcase_result(self,id:int,result:int,time:int,mem:int,status:JudgeResultCode):
        if id not in self.testcase_map: raise KeyError("Testcase not found")
        for i in self.testcase_map[id]:
            i.score=result
            i.time=time
            i.memory=mem
            i.status=status
    def _build_engine(self,now:ScorableObject,config:Union[SubtaskConfig,TestcaseConfig]):
        if "policy" not in config:
            if config["real_id"] not in self.testcase_map:
                self.testcase_map[config["real_id"]]=[]
            self.testcase_map[config["real_id"]].append(cast(Testcase,now))
            return
        else:
            if isinstance(now,Subtask):
                self.subtask_map[config["real_id"]]=now
                for i in config["sub"]:
                    obj:ScorableObject=None
                    if "policy" in i:
                        obj=now.add_subtask(i["policy"],i["real_id"])
                    elif "real_id" in i:
                        obj=now.add_testcase(0,0,0,i["real_id"])
                    else: raise InvalidConfigError()
                    self._build_engine(obj,i)
            else: raise InvalidConfigError()